27adc8
@@ -75,66 +75,26 @@
public class TopologySSEServlet extends HttpServlet {
         AsyncContext asyncContext = req.startAsync();
         PrintWriter writer = resp.getWriter();
 
-        Object writeLock = new Object();
-
-        TopologyListener topologyListener = topo -> {
-            String json = topologyToJson(req);
-            synchronized (writeLock) {
-                writer.write("event: topologyChange\n");
-                writer.write("data: " + json);
-                writer.flush();
-            }
-        };
-
-        ScheduledFuture keepAlive = this.keepAliveExecutor.scheduleAtFixedRate(() -> {
-                                                                                   try {
-                                                                                       writer.write(":\n\n");
-                                                                                       writer.flush();
-                                                                                   } catch (Throwable t) {
-                                                                                       TopologySSEServlet.this.topology.removeListener(topologyListener);
-                                                                                       throw t;
-                                                                                   }
-                                                                               },
-                                                                               10,
-                                                                               15,
-                                                                               TimeUnit.SECONDS);
-
+        TopologyListener topologyListener = new SSETopologyListener(writer, req.isSecure());
 
+        ScheduledFuture keepAlive = this.keepAliveExecutor.scheduleAtFixedRate(
+                new KeepAliveRunnable(writer, topologyListener),
+                10,
+                15,
+                TimeUnit.SECONDS);
         asyncContext.setTimeout(0);
-        asyncContext.addListener(new AsyncListener() {
-            @Override
-            public void onComplete(AsyncEvent asyncEvent) throws IOException {
-                TopologySSEServlet.this.topology.removeListener(topologyListener);
-                keepAlive.cancel(true);
-            }
-
-            @Override
-            public void onTimeout(AsyncEvent asyncEvent) throws IOException {
-                TopologySSEServlet.this.topology.removeListener(topologyListener);
-                keepAlive.cancel(true);
-            }
-
-            @Override
-            public void onError(AsyncEvent asyncEvent) throws IOException {
-                TopologySSEServlet.this.topology.removeListener(topologyListener);
-                keepAlive.cancel(true);
-            }
-
-            @Override
-            public void onStartAsync(AsyncEvent asyncEvent) throws IOException {
-            }
-        });
+        asyncContext.addListener(new TopologyAsyncListener(topology, topologyListener, keepAlive));
 
 
         this.topology.addListener(topologyListener);
-        String json = topologyToJson(req);
+        String json = topologyToJson(req.isSecure());
         writer.write("event: topologyChange\n");
         writer.write("data: " + json);
         writer.flush();
 
     }
 
-    protected String topologyToJson(HttpServletRequest req) {
+    private String topologyToJson(boolean secure) {
         StringBuilder json = new StringBuilder();
 
         json.append("{");
@@ -152,7 +112,7 @@
public class TopologySSEServlet extends HttpServlet {
             String proxyContext = getServletContext().getInitParameter(key + "-proxy");
             if (proxyContext != null) {
                 List<String> tags = new ArrayList<>();
-                tags.add(req.isSecure() ? "https" : "http");
+                tags.add(secure ? "https" : "http");
                 while (listIter.hasNext()) {
                     Topology.Entry server = listIter.next();
                     tags.add(server.getAddress() + ":" + server.getPort());
@@ -185,11 +145,10 @@
public class TopologySSEServlet extends HttpServlet {
         }
 
         json.append("}\n\n");
-
         return json.toString();
     }
 
-    protected void populateEndpointAndTagsJson(StringBuilder json, String endpoint, List<String> tags) {
+    private void populateEndpointAndTagsJson(StringBuilder json, String endpoint, List<String> tags) {
         json.append("{");
         json.append("\"endpoint\": \"").append(endpoint).append("\",");
         json.append("\"tags\":[");
@@ -208,4 +167,99 @@
public class TopologySSEServlet extends HttpServlet {
     private Topology topology;
 
     private ScheduledExecutorService keepAliveExecutor;
+
+    private class KeepAliveRunnable implements Runnable {
+        private final PrintWriter writer;
+
+        private final TopologyListener topologyListener;
+
+        public KeepAliveRunnable(PrintWriter writer, TopologyListener topologyListener) {
+            this.writer = writer;
+            this.topologyListener = topologyListener;
+        }
+
+        @Override
+        public void run() {
+            try {
+                writer.write(":\n\n");
+                writer.flush();
+            } catch (Throwable t) {
+                TopologySSEServlet.this.topology.removeListener(topologyListener);
+                throw t;
+            }
+
+        }
+    }
+
+    public static class TopologyAsyncListener implements AsyncListener {
+        private final ScheduledFuture keepAlive;
+
+        private final TopologyListener topologyListener;
+
+        private final Topology topology;
+
+        public TopologyAsyncListener() {
+            // Do Nothing
+            this.topology = null;
+            this.topologyListener = null;
+            this.keepAlive = null;
+        }
+
+        public TopologyAsyncListener(Topology topology, TopologyListener topologyListener, ScheduledFuture scheduledFuture) {
+            this.topology = topology;
+            this.topologyListener = topologyListener;
+            this.keepAlive = scheduledFuture;
+        }
+
+        @Override
+        public void onComplete(AsyncEvent asyncEvent) throws IOException {
+            if (topology != null) {
+                topology.removeListener(topologyListener);
+                keepAlive.cancel(true);
+            }
+        }
+
+        @Override
+        public void onTimeout(AsyncEvent asyncEvent) throws IOException {
+            if (topology != null) {
+                topology.removeListener(topologyListener);
+                keepAlive.cancel(true);
+            }
+        }
+
+        @Override
+        public void onError(AsyncEvent asyncEvent) throws IOException {
+            if (topology != null) {
+                topology.removeListener(topologyListener);
+                keepAlive.cancel(true);
+            }
+        }
+
+        @Override
+        public void onStartAsync(AsyncEvent asyncEvent) throws IOException {
+        }
+    }
+
+    private class SSETopologyListener implements TopologyListener {
+        Object writeLock = new Object();
+
+        final PrintWriter writer;
+
+        final boolean secure;
+
+        private SSETopologyListener(PrintWriter writer, boolean secure) {
+            this.writer = writer;
+            this.secure = secure;
+        }
+
+        @Override
+        public void onChange(Topology topology) {
+            String json = topologyToJson(secure);
+            synchronized (writeLock) {
+                writer.write("event: topologyChange\n");
+                writer.write("data: " + json);
+                writer.flush();
+            }
+        }
+    }
 }
